import numpy as np
from math import sin, cos, radians, sqrt
import matplotlib.pyplot as plt

# --- I. 物理常数 (除 rho_0 外) 和弹丸参数 ---
# 国际标准大气模型参数 (常量部分)
# RHO0 (海平面密度) 将由用户输入
T0 = 288.15          # 海平面温度 (K)
L = 0.0065           # 温度递减率 (K/m)
G = 9.80665          # 重力加速度 (m/s^2)
R = 287.05           # 气体常数 J/(kg*K) 
P_TERM = (G / (R * L)) - 1 # 密度公式指数项

# 弹丸参数 (假设值)
C_SOUND = 338.0      
M = 5.0              
A = 0.0095           

# Cd-Ma 关系数据
MA_POINTS = np.array([0.0, 0.7, 1.0, 1.3, 2.0, 2.5]) 
CD_POINTS = np.array([0.48, 0.45, 0.95, 0.45, 0.40, 0.40]) 

# --- II. 动态物理模型函数 (需传入 rho_0) ---

def get_drag_coefficient(v, c_sound):
    """根据马赫数线性插值计算 Cd。"""
    mach_number = v / c_sound
    # Cd 依赖于速度
    Cd = np.interp(mach_number, MA_POINTS, CD_POINTS) 
    return Cd

def calculate_initial_density(h0, rho0):
    """
    根据国际标准大气模型，计算初始海拔 h0 处的固定空气密度 $\rho(h_0)$。
    rho0 是用户输入的海平面密度。
    """
    if h0 < 0:
        return rho0
        
    temp_ratio = (T0 - L * h0) / T0
    if temp_ratio <= 0:
        return 0.0 
        
    # 国际标准大气模型公式，使用参数 rho0
    rho_h0 = rho0 * (temp_ratio) ** P_TERM 
    return rho_h0

# --- III. 弹道微分方程 (d S/d t) ---

def derivatives(S, params):
    """
    计算状态向量 S = [x, z, vx, vz] 对时间 t 的导数 dS/dt。
    params 中包含 rho_fixed。
    """
    x, z, vx, vz = S
    m, A, g, c_sound, rho_fixed = params 

    v = sqrt(vx**2 + vz**2) # 速度大小
    rho = rho_fixed # 密度固定在 rho(h0)
    
    # Cd 动态依赖于速度
    Cd = get_drag_coefficient(v, c_sound) 
    
    # 加速度计算
    drag_coeff = -(Cd * rho * A) / (2.0 * m) 
    
    dvx_dt = drag_coeff * v * vx 
    dvz_dt = drag_coeff * v * vz - g 
    
    return np.array([vx, vz, dvx_dt, dvz_dt])

# --- IV. RK4 积分函数 ---

def rk4_step(S, dt, params):
    """执行一步 RK4 积分。"""
    
    K1 = dt * derivatives(S, params)
    K2 = dt * derivatives(S + K1 / 2.0, params)
    K3 = dt * derivatives(S + K2 / 2.0, params)
    K4 = dt * derivatives(S + K3, params)
    
    S_new = S + (K1 + 2.0 * K2 + 2.0 * K3 + K4) / 6.0
    return S_new

# --- V. 弹道模拟核心函数 (返回射程 R) ---

def calculate_range_fixed_alt(h0, rho0, v0, theta_deg, dt=0.01):
    """
    使用 RK4 算法模拟弹道轨迹，固定初始海拔 h0 处的密度。
    需要 rho0 作为参数来计算 rho(h0)。
    """
    rho_fixed = calculate_initial_density(h0, rho0)
    # 将 rho_fixed 加入 params 传递给微分方程
    params = (M, A, G, C_SOUND, rho_fixed) 
    theta_rad = radians(theta_deg)
    
    # 初始状态 S = [x, z, vx, vz]
    S = np.array([
        0.0, 0.0, 
        v0 * cos(theta_rad), 
        v0 * sin(theta_rad)
    ])
    
    # 积分循环
    while S[1] >= -0.01: 
        S_new = rk4_step(S, dt, params)
        
        # 检查是否跨越地面
        if S_new[1] < 0 and S[1] >= 0:
            break 
            
        S = S_new

    R_final = S[0] # 返回 x 坐标作为射程
    
    return R_final, rho_fixed

# --- VI. 交互式运行和输出 ---

def interactive_run_fixed_alt_with_rho0_param(dt=0.01):
    
    print("\n--- 🎯 RK4 弹道计算器 (ρ₀为参数，固定初始海拔密度模型) ---")
    print("您需要输入海平面密度 $\\rho_0$，以及初始海拔 $h_0$。")
    
    try:
        # 获取用户输入
        rho0_input = float(input("请输入海平面密度 ρ₀ (kg/m³)："))
        h0 = float(input("请输入初始海拔 h₀ (m)："))
        v0 = float(input("请输入初速度 v₀ (m/s)："))
        theta = float(input("请输入入射角 θ (度, 0-90)："))
        
        if not (0 < theta <= 90) or v0 <= 0 or h0 < 0 or rho0_input <= 0:
            print("输入参数无效。请确保 ρ₀ > 0, h₀ >= 0, v₀ > 0 且 0 < θ <= 90。")
            return

    except ValueError:
        print("输入错误。请确保输入的是数字。")
        return
    except Exception as e:
        print(f"发生错误: {e}")
        return

    # --- 计算和输出 ---
    print(f"\n--- 📈 计算开始 (dt={dt}s) ---")
    
    # 1. 计算固定密度值 (依赖于用户输入的 rho0_input)
    R_final, rho_fixed_value = calculate_range_fixed_alt(
        h0, rho0_input, v0, theta, dt=dt
    )
    
    print("\n--- 🏁 模拟结果 ---")
    print(f"设定参数: 海平面密度 $\\rho_0$ = {rho0_input:.4f} kg/m³")
    print(f"设定参数: 初始海拔 h₀ = {h0:.1f} m")
    print(f"   --> 固定弹道密度 $\\rho(h_0)$ = {rho_fixed_value:.4f} kg/m³")
    print(f"设定参数: v₀ = {v0:.3f} m/s, θ = {theta:.1f}°")
    print(f"最终射程 R: {R_final:.3f} m")

# 启动交互式运行
if __name__ == "__main__":
    interactive_run_fixed_alt_with_rho0_param()